#ifndef NN_MATCHING_H
#define NN_MATCHING_H

//#include <torch/torch.h>

#include <map>
#include <vector>

#include <Eigen/Core>
#include <Eigen/Dense>
#include <opencv2/opencv.hpp>

std::vector<std::vector<float>> iou_dist(const std::vector<cv::Rect2f> &dets, const std::vector<cv::Rect2f> &trks);

// save features of the track in GPU
class FeatureBundle {
public:
    FeatureBundle() : full(false), next(0){
    }

    void clear() {
        next = 0;
        full = false;
    }

    bool empty() const {
        return next == 0 && !full;
    }

    void add(std::vector<float> &feat) {
        if (next == budget) {
            full = true;
            next = 0;
        }
//        assert(feat.size()==feat_dim);
        store.emplace_back(feat);
        next++;
    }

    std::vector<std::vector<float>> get() const {
        return store;
    }

private:
    // budget is lower, speed up
    static const int64_t budget = 48, feat_dim = 512;

    std::vector<std::vector<float>> store;

    bool full;
    int64_t next;
};

template<typename TrackData>
class FeatureMetric {
public:
    explicit FeatureMetric(std::vector<TrackData> &data) : data(data) {}

    std::vector<std::vector<float>>  distance(const std::vector<std::vector<float>> &features, const std::vector<int> &targets) {
        std::vector<std::vector<float>> dist{};
        if (features.size()) {
            //std::cout << "############# " << features.size() << std::endl;
            for (size_t i = 0; i < targets.size(); ++i) {
                dist.emplace_back(nn_cosine_distance(data[targets[i]].feats.get(), features));
            }
        }

        return dist;
    }

    void update( std::vector<std::vector<float>> &feats, const std::vector<int> &targets) {

        for (size_t i = 0; i < targets.size(); ++i) {
            data[targets[i]].feats.add(feats[i]);
        }
    }

private:
    std::vector<TrackData> &data;
    typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> FEATURESS;
    Eigen::MatrixXf _cosine_distance(const FEATURESS & a, const FEATURESS & b)
    {
        FEATURESS aa = a;
        FEATURESS bb = b;
        for (int i = 0; i < a.rows(); ++i)
        {
            aa.row(i) =  a.row(i) / sqrt(a.row(i).squaredNorm());
        }
        for (int i = 0; i < b.rows(); ++i)
        {
            bb.row(i) =  b.row(i) / sqrt(b.row(i).squaredNorm());
        }

        Eigen::MatrixXf res = 1. - (aa * bb.transpose()).array();

        return res;
    }

    std::vector<float> nn_cosine_distance(std::vector<std::vector<float>> feats, std::vector<std::vector<float>> features)
    {
        float * feats_ = (float *)malloc(feats.size() * feats[0].size() * sizeof(float));
        float * features_ = (float *)malloc(features.size() * features[0].size() * sizeof(float));
        int offset_bin_1=feats[0].size()*sizeof(float);
        int offset_bin_2=features[0].size()*sizeof(float);
        for (size_t i =0;i<feats.size(); i++){
            memcpy((void *)(feats_)+i*offset_bin_1, (void *)feats[i].data(), offset_bin_1);
        }
        for (size_t i =0;i<features.size(); i++){

            memcpy((void *)(features_)+i*offset_bin_2, (void *)features[i].data(), offset_bin_2);
        }
        FEATURESS track = Eigen::Map<FEATURESS>(feats_, feats.size(), feats[0].size());
        FEATURESS det = Eigen::Map<FEATURESS>(features_, features.size(), features[0].size());
        Eigen::MatrixXf distances = _cosine_distance(track, det);
        delete [] feats_, features_;
        Eigen::VectorXf res = distances.colwise().minCoeff();

        return std::vector<float>{res.data(), res.data()+res.size()};
    }
    // is shit
    // std::vector<float> nn_cosine_distance(std::vector<std::vector<float>> feats, std::vector<std::vector<float>> features) {
    //     cv::Mat feats_ = cv::Mat(feats.size(), feats[0].size(), CV_32FC1, feats.data());
    //     cv::Mat features_ = cv::Mat(features.size(), features[0].size(), CV_32FC1, features.data());
    //     cv::Mat cosine_similarity = (1 - feats_ * features_.t()).t();
    //     std::vector<float> res;
    //     for (size_t i=0; i<cosine_similarity.rows-1; i++){
    //         cv::Mat row = cosine_similarity.colRange(i , i+1);
    //         double maxvalue=0, minvalue=0;
    //         cv::minMaxLoc(row, &minvalue, &maxvalue);
    //         res.emplace_back(float(minvalue));
    //     }
    //     return res;
    // }
};

#endif //NN_MATCHING_H
